#!/usr/bin/env python

import pygame
from pygame.locals import *

import random as rand

import helper
import image
import rule
import tile

from helper import *
from image import *
from tile import draw_tile

CLASS_TREADS = 'u_treads'
CLASS_FOOT = 'u_foot'
CLASS_WINGS = 'u_wings'
CLASS_WATER = 'u_water'

CURSOR_NORMAL = 'cursor_normal'
CURSOR_ATTACK = 'cursor_attack'

ARR_LT  = 'left_top'
ARR_LR  = 'left_right'
ARR_TB  = 'top_bottom'
ARR_LB  = 'left_bottom'
ARR_RT  = 'right_top'
ARR_RB  = 'right_bottom'
ARR_L   = 'left'
ARR_R   = 'right'
ARR_T   = 'top'
ARR_B   = 'bottom'

def load_arrow(arrow_basename, *args, **properties):
  sc = get_prop(properties, 'scale', (1.0, 1.0))
  sz = get_prop(properties, 'size', (100, 100))
  al = get_prop(properties, 'alpha', 1.0)
  
  arrows = { }
  for dir in [ARR_LT, ARR_LR, ARR_TB, ARR_LB, ARR_RT, ARR_RB, ARR_L,
              ARR_R, ARR_T, ARR_B]:
    arrows[dir] = Still(svg_file_to_surface(file('arrow_%s.svg' % (dir)), size=sz, scale=sc, alpha=al))
  return arrows

def load_cursor(*args, **properties):
  cursor = { }
  cursor[CURSOR_NORMAL] = get_prop(properties, CURSOR_NORMAL, must_exist=True)
  cursor[CURSOR_ATTACK] = get_prop(properties, CURSOR_ATTACK, def_val=cursor[CURSOR_NORMAL])
  return cursor






class UnitError(Exception):
  def __init__(self, value):
    self.value = value
  
  def __str__(self):
    return repr(self.value)

STATE_NORMAL = 'state_normal'
STATE_SELECTED = 'state_selected'
STATE_MOVING = 'state_moving'
STATE_MOVED = 'state_moved'
STATE_DESTROYED = 'state_destroyed'

MOV_INF = 10000

MAX_UNITS = 6000
MAX_UNIT_NAME_LEN = 6

class Unit:
  def __init__(self, unit_type, player=None):
    self.player =         player
    self.unit_type =      unit_type
    self.unit_name =      unit_type.get_name() #maybe different if want unique names
    
    self.properties =     { }
    
    self.properties['hp'] =             self.init_prop('max_hp')
    self.properties['mp'] =             self.init_prop('max_mp') #movement points
    self.properties['fuel'] =           self.init_prop('max_fuel')
    self.properties['state'] =          STATE_NORMAL
    
    
    self.id = -1
  
  def get_type(self):
    return self.unit_type
  
  def set_id(self, id):
    self.id = id
  
  def get_id(self):
    return self.id
  
  def get_state(self):
    return self.properties['state']
  
  def set_state(self, state):
    self.properties['state'] = state
  
  def change_fuel(self, amt):
    #if self.properties['fuel'] - amt < 0: self.properties['fuel'] = 0
    #else: self.properties['fuel'] += amt
    max_fuel = self.get_type().get_prop('max_fuel')
    
    chged = self.properties['fuel'] + amt
    if chged > max_fuel:
      chged = max_fuel
    elif chged < 0:
      chged = 0
    
    self.properties['fuel'] = chged
      
    
  
  def get_property(self, prop_name):
    return get_prop(self.properties, prop_name)
  
  def get_unit_type_property(self, prop_name):
    return self.unit_type.get_prop(prop_name)
 
  #If the must_exist flag is set, this function will complain
  #if the property isn't found
  def init_prop(self, prop_name, must_exist=True):
    
    val = self.unit_type.get_prop(prop_name)
    if val is None and must_exist:
      raise UnitError('Property %s is required but does not exist in unit type %s.' \
                      % (val, self.unit_type.get_name()))
    
    return val
  
  def update(self):
    anim = self.get_unit_type_property(self.properties['state'])
    if anim is None: anim = self.get_unit_type_property(STATE_NORMAL)
    return anim.update()
  
  def __eq__(self, other):
    return self.id == other.get_id()
  
  def __ne__(self, other):
    return not self.__eq__(other)
  



class UnitType:
  def __init__(self, name, unit_class, **properties):
    self.name = name
    self.properties = properties
    
    if STATE_NORMAL not in self.properties:
      raise UnitError('No normal state animation.')
    
    self.short_name = helper.get_prop(properties, 'short_name', self.name[0:MAX_UNIT_NAME_LEN])
    if len(self.short_name) > MAX_UNIT_NAME_LEN:
      raise UnitError('Unit short name cannot be longer than %d characters.' % (MAX_UNIT_NAME_LEN))
    
    self.set_unit_class(unit_class)
    self.set_unit_moved_anim()
  
  def set_unit_class(self, cl):
    if cl not in [CLASS_TREADS, CLASS_FOOT, CLASS_WATER, CLASS_WINGS]:
      raise UnitError('Invalid class %s.' % (cl))
    self.unit_class = cl
  
  def get_unit_class(self):
    return self.unit_class
  
  def set_unit_moved_anim(self):
    norm_anim = self.properties[STATE_NORMAL]
    self.properties[STATE_MOVED] = image.Grayscale(norm_anim.update())
    pass
  
  def get_name(self):
    return self.name
  
  def get_short_name(self):
    return self.short_name
  
  def get_anim(self, state):
    return self.get_prop(state)
  
  def get_prop(self, prop_name):
    return helper.get_prop(self.properties, prop_name)
  
  #each unit type has a unique name, class and set of properties
  def __eq__(self, other):
    return self.properties == other.properties and \
        self.name == other.get_name() and \
        self.unit_class == other.get_unit_class()
  
  def __ne__(self, other):
    return not self.__eq__(other)


class UnitEngine:
  #init the tile engine first before calling this
  def __init__(self, tile_engine, rule_engine):
    self.tile_engine = tile_engine
    self.rule_engine = rule_engine
    map_cols, map_rows = tile_engine.get_map_dimensions()
    
    self.map_dimen = (map_cols, map_rows)
    self.map2d = [None] * map_rows
    for i in range(0, map_rows):
      self.map2d[i] = [None] * map_cols

    #a dictionary of long names
    self.unit_types = { }
    self.units = { }
    
    self.selected_unit_id = None
    self.selected_unit_m_region = []
    self.selected_unit_path = []
    self.selected_unit_is_moving = False
    self.selected_unit_curr_path_index = None
    self.selected_unit_real_pos = None
    
    #TODO improve the algorithm for drawing a moving unit, so that
    #situations where tile dimensions that are not a multiple of the 
    #speed are handled correctly
    t_wh = self.tile_engine.get_tile_wh()
    self.selected_unit_mv_xspeed = t_wh[0] / 5
    self.selected_unit_mv_yspeed = t_wh[1] / 5
    
    self.mv_highlight_tile = None
    self.arrow = None
    
    self.cursor = None
    self.cursor_state = CURSOR_NORMAL
    self.cursor_grid_xy = (0, 0)
  
  def set_arrow(self, arrow):
    self.arrow = arrow
  
  def set_cursor(self, cursor):
    self.cursor = cursor
  
  def set_mv_highlight_tile(self, tile):
    self.mv_highlight_tile = tile
  
  def register_unit_type(self, unit_type):
    self.unit_types[unit_type.get_name()] = unit_type
  
  def place_unit(self, unit, grid_xy):
    id = unit.get_id()
    if id in self.units:
      raise UnitError('Unit already exists in the engine.')
    elif id == -1:
      id = gen_id(0, MAX_UNITS * 1000 - 1, self.units.keys())
    unit.set_id(id)
    
    type = unit.get_type()
    if type.get_name() not in self.unit_types:
      self.register_unit_type(type)
    
    tile_type = self.tile_engine.get_tile_type_on_grid(grid_xy)
    p_rule = self.rule_engine.get_placement_rule(type, tile_type)
    if p_rule is not None:
      can_place = p_rule.get_property('can_place')
      if not can_place:
        raise UnitError('Cannot place unit %s on tile %s.' % (type.get_name(), tile_type.get_name()))
    
    grid_x, grid_y = grid_xy
    if self.map2d[grid_y][grid_x] is not None:
      raise UnitError('Tile (%d, %d) is already occupied.' % (grid_x, grid_y))
    self.set_unit_on_grid(unit, grid_xy)
  
  def update_unit(self, unit):
    #TODO tell the unit engine to update unit status:
    # - the unit animation state (STATE_NORMAL)
    # - the unit hp/fuel (if it is on a base/airfield/whatever)
    unit.set_state(STATE_NORMAL)
    grid_xy = self.units[unit.get_id()][1]
    tile_type = self.get_tile_type_on_grid(grid_xy)
    
    h_rule = self.rule_engine.get_hp_rule(unit.get_type(), tile_type)
    f_rule = self.rule_engine.get_fuel_rule(unit.get_type(), tile_type)
    
    if h_rule is not None:
      pass
    
    pass
  
  def get_tile_wh(self):
    return self.tile_engine.get_tile_wh()
  
  def real_to_grid(self, r_xy):
    return self.tile_engine.real_to_grid(r_xy)
  
  def get_tile_type_on_grid(self, grid_xy):
    return self.tile_engine.get_tile_type_on_grid(grid_xy)
  
  def set_unit_on_grid(self, unit, grid_xy):
    #remove the unit from the grid if it was originally there
    if unit.get_id() in self.units:
      u, curr_xy = self.units[unit.get_id()]
      self.map2d[curr_xy[1]][curr_xy[0]] = None
    
    grid_x, grid_y = grid_xy
    self.map2d[grid_y][grid_x] = unit
    self.units[unit.get_id()] = (unit, grid_xy)
  
  def get_unit_on_grid(self, grid_xy):
    grid_x, grid_y = grid_xy
    return self.map2d[grid_y][grid_x] #can be None
  
  def get_unit_move_cost_of_path(self, unit, path):
    unit_type = unit.get_type()
    cost = 0
    for grid_xy in path:
      tile_type = self.get_tile_type_on_grid(grid_xy)
      cost += self.get_unit_move_cost(unit_type, tile_type)
    return cost
  
  def get_unit_move_cost(self, unit_type, tile_type):
    #cost = None
    m_rule = self.rule_engine.get_movement_rule(unit_type, tile_type)
    if m_rule is None: return 1
    else:
      return m_rule.get_property('cost')
      #return cost
  
  def get_unit_movable_region(self, unit_id):
    unit, grid_xy = self.units[unit_id]
    unit_type = unit.get_type()
    mp = unit.get_property('mp')
    fuel = unit.get_property('fuel')
    if fuel > mp: start_mp = mp
    else: start_mp = fuel
    
    dirs = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    
    visited = set()
    queue = [(grid_xy, start_mp)]
    
    def in_queue(xy):
      for g_xy, mp in queue:
        if g_xy == xy:
          return True
      return False
    
    while len(queue) != 0:
      curr_xy, rem_mp = queue.pop(0)
      visited.add(curr_xy)
      if rem_mp == 0:
        continue
      for dx, dy in dirs:
        new_xy = (curr_xy[0] + dx, curr_xy[1] + dy)
        if new_xy[0] < 0 or new_xy[1] < 0: continue
        if new_xy[0] >= self.map_dimen[0] or new_xy[1] >= self.map_dimen[1]: continue
        if new_xy in visited or in_queue(new_xy): continue
        
        tt = self.get_tile_type_on_grid(new_xy)
        cost = self.get_unit_move_cost(unit_type, tt)
        if cost == MOV_INF or rem_mp - cost < 0: continue
        
        queue.append((new_xy, rem_mp - cost))
    
    return visited # a list of 2-tuples representing grid coordinates
  
  def is_unit_selected(self):
    return self.selected_unit_id is not None
  
  def select_unit(self, unit_id):
    #TODO check if the player can select the unit, and in what way
    sel_unit, xy = self.units[unit_id]
    if sel_unit.get_state() == STATE_MOVED:
      return
    self.selected_unit_id = unit_id
    sel_unit.set_state(STATE_SELECTED)
    if len(self.selected_unit_m_region) == 0:
      self.selected_unit_m_region = self.get_unit_movable_region(unit_id)
    self.selected_unit_curr_path_index = 0
    t_wh = self.tile_engine.get_tile_wh()
    #the 'real position' is always at the center of the tile
    self.selected_unit_real_pos = [xy[0] * t_wh[0], xy[1] * t_wh[1]]
  
  def deselect_unit(self, unit_id):
    self.selected_unit_id = None
    self.selected_unit_m_region = []
    self.selected_unit_path = []
    sel_unit, xy = self.units[unit_id]
    sel_unit.set_state(STATE_NORMAL)
    self.selected_unit_curr_path_index = None
    self.selected_unit_real_pos = None
  
  def calc_path(self, unit, grid_xy1, grid_xy2):
    def sgn(x):
      if x > 0: return 1
      elif x == 0: return 0
      else: return -1
    
    def manhattan(g1, g2):
      g1_x, g1_y = g1
      g2_x, g2_y = g2
      xdiff = (g1_x - g2_x) * sgn(g1_x - g2_x)
      ydiff = (g1_y - g2_y) * sgn(g1_y - g2_y)
      return xdiff + ydiff
    
    #TODO use A*
    closed_set = set()
    open_set = set()
    open_set.add(grid_xy1)
    
    g_score = { grid_xy1: 0 }
    h_score = { grid_xy1: manhattan(grid_xy1, grid_xy2) }
    f_score = { grid_xy1: h_score[grid_xy1]}
    came_from = { }
    
    def next():
      return min(open_set, key=lambda x: f_score[x])
    
    def reconstruct_path(curr_xy):
      path = [curr_xy]
      while curr_xy in came_from:
        path.insert(0, came_from[curr_xy])
        curr_xy = came_from[curr_xy]
      return path
    
    neighbours = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    
    while len(open_set) != 0:
      n = next()
      if n == grid_xy2: #GOL!!!
        return reconstruct_path(grid_xy2)
      open_set.remove(n)
      closed_set.add(n)
      for nn in neighbours:
        new_xy = (n[0] + nn[0], n[1] + nn[1])
        if new_xy[0] < 0 or new_xy[1] < 0: continue
        if new_xy[0] >= self.map_dimen[0] or new_xy[1] >= self.map_dimen[1]: continue
        if new_xy in closed_set: continue
        
        tt = self.get_tile_type_on_grid(new_xy)
        move_cost = self.get_unit_move_cost(unit.get_type(), tt)
        if move_cost == MOV_INF: continue
        
        tent_g_score = g_score[n] + move_cost
        tent_is_better = False
        if new_xy not in open_set:
          open_set.add(new_xy)
          h_score[new_xy] = manhattan(new_xy, grid_xy2)
          tent_is_better = True
        elif tent_g_score < g_score[new_xy]:
          tent_is_better = True
        
        if tent_is_better:
          came_from[new_xy] = n
          g_score[new_xy] = tent_g_score
          f_score[new_xy] = g_score[new_xy] + h_score[new_xy]
    return [] # a list of 2-tuples representing grid coordinates
  
  def handle_cursor_change(self, grid_xy):
    self.set_cursor_position(grid_xy)
    
    if self.selected_unit_is_moving: return
    if self.selected_unit_id is None:
      return
    else:
      #calculate the path from the unit's xy to the cursor's xy (grid_xy)
      if grid_xy in self.selected_unit_m_region:
        sel_unit, xy = self.units[self.selected_unit_id]
        self.selected_unit_path = self.calc_path(sel_unit, xy, grid_xy)
  
  def handle_cursor_select(self, grid_xy):
    self.set_cursor_position(grid_xy)
    
    if self.selected_unit_is_moving: return
    #TODO center the viewport on the grid_xy
    unit = self.get_unit_on_grid(grid_xy)
    if unit is not None:
      if self.selected_unit_id is None: self.select_unit(unit.get_id())
      elif unit.get_id() == self.selected_unit_id: self.deselect_unit(unit.get_id())
      else:
        self.deselect_unit(self.selected_unit_id)
        self.select_unit(unit.get_id())
    else:
      if self.selected_unit_id is None: return
      else:
        #check if the clicked grid coord is within the unit's movable region
        if grid_xy not in self.selected_unit_m_region: return
        self.selected_unit_is_moving = True
        sel_unit, xy = self.units[self.selected_unit_id]
        sel_unit.set_state(STATE_MOVING)
        pass #TODO move the unit
  
  def set_cursor_position(self, grid_xy):
    self.cursor_grid_xy = grid_xy
  
  def get_path_index(self, curr_g_xy):
    for i in range(0, len(self.selected_unit_path)):
      if self.selected_unit_path[i] == curr_g_xy:
        return i
    return -1
  
  def move_unit(self, xspeed=None, yspeed=None):
    if xspeed is None: xspeed = self.selected_unit_mv_xspeed
    if yspeed is None: yspeed = self.selected_unit_mv_yspeed
    
    if self.selected_unit_is_moving:
      t_wh = self.get_tile_wh()
      to_grid_xy = self.selected_unit_path[self.selected_unit_curr_path_index]
      to_real_xy = [to_grid_xy[0] * t_wh[0], to_grid_xy[1] * t_wh[1]]
      curr_real_xy = self.selected_unit_real_pos
      #print to_real_xy, curr_real_xy
      if to_real_xy == curr_real_xy:
        self.selected_unit_curr_path_index += 1
        if self.selected_unit_curr_path_index == len(self.selected_unit_path):
          
          self.selected_unit_is_moving = False
          sel_unit, xy = self.units[self.selected_unit_id]
          
          path_length = len(self.selected_unit_path)
          path_cost = self.get_unit_move_cost_of_path(sel_unit, self.selected_unit_path)
          
          self.set_unit_on_grid(sel_unit, self.selected_unit_path[len(self.selected_unit_path) - 1])
          self.deselect_unit(self.selected_unit_id)
          
          sel_unit.set_state(STATE_MOVED)
          sel_unit.change_fuel(-path_cost)
          return
        to_grid_xy = self.selected_unit_path[self.selected_unit_curr_path_index]
        to_real_xy = (to_grid_xy[0] * t_wh[0], to_grid_xy[1] * t_wh[1])
      
      if to_real_xy[0] < curr_real_xy[0]:
        self.selected_unit_real_pos[0] -= xspeed
      elif to_real_xy[0] > curr_real_xy[0]:
        self.selected_unit_real_pos[0] += xspeed
      elif to_real_xy[1] < curr_real_xy[1]:
        self.selected_unit_real_pos[1] -= yspeed
      elif to_real_xy[1] > curr_real_xy[1]:
        self.selected_unit_real_pos[1] += yspeed
      
  
  def draw_moving_unit(self):
    if self.selected_unit_is_moving:
      sel_unit, xy = self.units[self.selected_unit_id]
      t_wh = self.get_tile_wh()
      x = self.selected_unit_real_pos[0] + t_wh[0] / 2
      y = self.selected_unit_real_pos[1] + t_wh[1] / 2
      unit_upd = sel_unit.update()
      upd_rect = unit_upd.get_rect()
      map = self.tile_engine.update()
      
      #print self.selected_unit_path
      i = self.get_path_index(self.real_to_grid(self.selected_unit_real_pos))
      #print i, self.selected_unit_real_pos, (x, y), self.real_to_grid((x, y))
      
      dirty_sqs = []
      if i != 0:
        dirty_sqs.append(self.selected_unit_path[i - 1])
      dirty_sqs.append(self.selected_unit_path[i])
      if i != len(self.selected_unit_path) - 1:
        dirty_sqs.append(self.selected_unit_path[i + 1])
      
      for d in dirty_sqs:
        self.tile_engine.set_dirty(d)
      
      off_x = t_wh[0] - upd_rect.w
      off_y = t_wh[1] - upd_rect.h
      
      map.blit(unit_upd, (self.selected_unit_real_pos[0] + off_x / 2, 
                          self.selected_unit_real_pos[1] + off_y / 2))
      
    pass
  
  def draw_on_tile(self, surf, grid_xy):
    return self.tile_engine.draw_on_tile(surf, grid_xy)
  
  def draw_cursor(self):
    #self.cursor_state = CURSOR_ATTACK
    self.draw_on_tile(self.cursor[self.cursor_state].update(), 
                      self.cursor_grid_xy)
  
  def draw_arrow(self, path):
    def to_left(p1, p2): #p1 to the left of p2
      return p1[0] == p2[0] - 1
    def to_right(p1, p2):
      return p1[0] == p2[0] + 1
    def to_above(p1, p2):
      return p1[1] == p2[1] - 1
    def to_below(p1, p2):
      return p1[1] == p2[1] + 1
    
    if self.arrow is None:
      return
    if len(path) <= 1: return
    
    anim = None
    
    for i in range(0, len(path)):
      if i == len(path) - 1:
        p_behind, p= path[i - 1], path[i]
        if to_left(p_behind, p):
          anim = ARR_R
        elif to_right(p_behind, p):
          anim = ARR_L
        elif to_above(p_behind, p):
          anim = ARR_B
        elif to_below(p_behind, p):
          anim = ARR_T
      elif i == 0: continue #TODO create the arrow stubs
      else:
        p_behind, p, p_in_front = path[i - 1], path[i], path[i + 1]
        if to_left(p_behind, p) and to_left(p, p_in_front):
          anim = ARR_LR
        elif to_right(p_behind, p) and to_right(p, p_in_front):
          anim = ARR_LR
        elif to_above(p_behind, p) and to_above(p, p_in_front):
          anim = ARR_TB
        elif to_below(p_behind, p) and to_below(p, p_in_front):
          anim = ARR_TB
        
        elif to_left(p_behind, p) and to_below(p, p_in_front):
          anim = ARR_LT
        elif to_above(p_behind, p) and to_right(p, p_in_front):
          anim = ARR_LT
        elif to_left(p_behind, p) and to_above(p, p_in_front):
          anim = ARR_LB
        elif to_below(p_behind, p) and to_right(p, p_in_front):
          anim = ARR_LB
        
        elif to_right(p_behind, p) and to_below(p, p_in_front):
          anim = ARR_RT
        elif to_above(p_behind, p) and to_left(p, p_in_front):
          anim = ARR_RT
        elif to_right(p_behind, p) and to_above(p, p_in_front):
          anim = ARR_RB
        elif to_below(p_behind, p) and to_left(p, p_in_front):
          anim = ARR_RB
      if anim is not None:
        self.draw_on_tile(self.arrow[anim].update(), p)
      
    pass
  
  def update(self):
    #TODO add unit sprites, selected rects, arrows etc after getting the tile map
    t_map = self.tile_engine.update()
    t_map_rect = t_map.get_rect()
    
    if not self.selected_unit_is_moving:
      if self.mv_highlight_tile is not None:
        h_tile = self.mv_highlight_tile.update()
        for g_x, g_y in self.selected_unit_m_region:
          self.draw_on_tile(h_tile, (g_x, g_y))
    
    self.move_unit()
    self.draw_moving_unit()
    
    for id in self.units:
      if id == self.selected_unit_id and self.selected_unit_is_moving: continue
      unit, xy = self.units[id]
      #draw_tile(unit.update(), self.map, t_wh, xy)
      self.draw_on_tile(unit.update(), xy)
    
    if not self.selected_unit_is_moving:
      self.draw_arrow(self.selected_unit_path)
    
    if self.cursor is not None:
      self.draw_cursor()
    
    return t_map
